from typing import List

import numpy as np

from spatialmath import SE3

from .. import Geometry
from ..simplex import Point, UnitVector, Support
from .geometry3d import Geometry3D


class Sphere(Geometry3D, Support):
    def __init__(self, base: SE3, radius: float) -> None:
        super().__init__(base)
        self.radius = radius

    @property
    def points(self) -> List[Geometry]:
        return []

    def calculate_support_point(self, d: UnitVector) -> Point:
        return Point(self.get_t() + self.radius * d.get_t())

    def plot(self, ax):
        phi, theta = np.mgrid[0:np.pi:20j, 0:2 * np.pi:20j]
        x = self.radius * np.cos(theta) * np.sin(phi)
        y = self.radius * np.sin(theta) * np.sin(phi)
        z = self.radius * np.cos(phi)

        coordinates = self.__calculate_coordinates(x, y, z)
        ax.plot_surface(*coordinates, color='b', alpha=0.5)

    def __calculate_coordinates(self, x: np.ndarray, y: np.ndarray, z: np.ndarray):
        x_new = np.zeros_like(x)
        y_new = np.zeros_like(y)
        z_new = np.zeros_like(z)

        num1, num2 = x.shape

        for i in range(num1):
            for j in range(num2):
                t = np.squeeze(self.base * [x[i, j], y[i, j], z[i, j]])
                [x_new[i, j], y_new[i, j], z_new[i, j]] = t
        return x_new, y_new, z_new